# coding: utf-8
"""
@Time    : 2024/8/28 10:14
@Author  : Y.H LEE
"""
from sys_params import device
import numpy as np
import torch.nn

from benchmarks.test_algorithms.genetic_algorithm import GeneticAlgorithm
from benchmarks.test_algorithms.particle_swarm_optimization import ParticleSwarmOptimization

from tqdm import tqdm

from utils.tools import *

'''
Trainer
'''


class GridTrainer:

    def __init__(self):
        self.dist = []

    def __call__(self, h_params, metric_dict, env, save_model=False, early_stop=False):
        """build algorithm"""
        algorithm = self.__get_algorithm__(h_params)

        """start training & evaluating"""
        train_logs = h_params['train_logs_save_dir'] + h_params['train_logs']
        with open(train_logs, 'w', encoding='utf-8') as fw:
            for k, v in h_params.items():
                fw.write('--> ' + k + ' : ' + str(v) + '\n')
            fw.write('\n')

        # for loop in tqdm(range(h_params['LOOPS']), total=h_params['LOOPS'], position=0):
            # print('loop:{0}'.format(loop))

        """algorithm training"""
        dist = algorithm.train(env, h_params['LOOPS'])
        self.dist = dist

        best = dist[-1]

        with open(train_logs, 'a', encoding='utf-8') as fw:
            fw.write(f'Best performing : \n'
                     f'      min_distance = {round(best, 4)}\n')

        # plot training metric graph
        plt.figure(2, figsize=(10, 5))
        plt.title(" Training metric")
        plt.plot(self.dist, label="distance")
        plt.xlabel("time step (=iterations)")
        plt.ylabel("distance")
        plt.legend()
        plt.savefig(h_params['train_logs_save_dir'] + '/train.png')
        plt.clf()

        print(f'min_distance: {best}')

        return best

    def __get_algorithm__(self, h_params):
        algorithm_name = h_params['ALGORITHM']
        algorithm = None
        if algorithm_name == 'genetic_algorithm':
            algorithm = GeneticAlgorithm(h_params["N_CITIES"], h_params["CROSS_RATE"], h_params["MUTATE_RATE"],
                                         h_params["POP_SIZE"])
        elif algorithm_name == 'particle_swarm_optimization':
            algorithm = ParticleSwarmOptimization(h_params["N_CITIES"], h_params["R1"], h_params["R2"],
                                                  h_params["POP_SIZE"], )
        else:
            raise ValueError(f'algorithm {algorithm_name} not exists')

        return algorithm

    def reset_res_list(self):
        self.dist = []
